/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * Global Trust Authority is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{collections::HashMap, error::Error};
use serde_cbor::{self};
use plugin_manager::{AgentHostFunctions, AgentPlugin, CollectEvidenceParams, PluginBase, PluginError, QueryConfigurationFn};

use crate::{ascend_npu_sdk::{ASCEND_NPU_API_CONFIG, AttestCtx, EVIDENCE_TYPE_EVIDENCE_TYPE_BOOT_MEASUREMENTS, EVIDENCE_TYPE_EVIDENCE_TYPE_QUOTE, EVIDENCE_TYPE_EVIDENCE_TYPE_RUNTIME_MEASUREMENTS}, config::AscendNpuConfig, entity::DcmiEvidence};
use crate::ascend_npu_sdk::init_ascend_npu_api;
use crate::config::{ASCEND_NPU_IMA_LOG, ASCEND_NPU_TCG_EVENT_LOG};
use crate::entity::AscendNpuEvidence;

const ASCEND_NPU_PLUGIN_TYPE: &str = "ascend_npu";

const OUT_OF_MEMORY: i32 = 0x2F;

#[derive(Debug)]
pub struct AscendNpuPlugin {
    pub config: AscendNpuConfig,
}

impl AscendNpuPlugin {
    pub fn new(plugin_type: String, query_configuration: QueryConfigurationFn) -> Result<Self, PluginError> {
        if plugin_type != ASCEND_NPU_PLUGIN_TYPE {
            return Err(PluginError::InputError(format!(
                "Plugin type {} not match {} plugin type",
                plugin_type, ASCEND_NPU_PLUGIN_TYPE
            )));
        }
        let config_json = query_configuration(plugin_type.clone()).ok_or_else(|| {
            PluginError::InputError(format!("Failed to query configuration for plugin type {}", plugin_type))
        })?;
        let config = AscendNpuConfig::from_json(plugin_type, &config_json)?;
        unsafe {
            match init_ascend_npu_api(&config.ascend_sdk_path) {
                Ok(_) => {},
                Err(e) => return Err(PluginError::InternalError(format!("Failed to init ascend_npu config: {}", e))),
            }
        }
        Ok(Self { config })
    }
}

impl PluginBase for AscendNpuPlugin {
    fn plugin_type(&self) -> &str {
        ASCEND_NPU_PLUGIN_TYPE
    }
}

impl AgentPlugin for AscendNpuPlugin {
    /// # collect_evidence
    ///
    /// Collects evidence from the Ascend NPU device.
    ///
    /// # Parameters
    /// - `node_id`: Optional node ID.
    /// - `nonce`: Optional nonce for challenge-response.
    /// - `log_types`: Optional vector of log types to collect.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// - The Ascend NPU API configuration is not initialized.
    /// - Failed to collect the specified log types.
    ///
    /// # Returns
    /// A `Result` containing the serialized JSON evidence or a `PluginError`.
    /// let params = CollectEvidenceParams{
    //         node_id: Some(node_id.clone()),
    //         nonce: nonce_bytes.clone(),
    //         log_types: info.log_types.clone(),
    //         attester_ids: info.attester_ids.clone(),
    //     };
    fn collect_evidence(&self, evidence_params: CollectEvidenceParams) -> Result<serde_json::Value, PluginError> {
        let mut challenge = match evidence_params.nonce {
            Some(n) => n.to_vec(),
            None => vec![0u8; 32],
        };
        let evidence = self.collect_log(evidence_params.log_types, &mut challenge)?;
        let result = serde_json::to_value(evidence).map_err(|e| PluginError::InternalError(format!("Failed to serialize evidence: {}", e)));
        Ok(result?)
    }
}

impl AscendNpuPlugin {
    /// # collect_log
    ///
    /// Collects log types from the Ascend NPU device.
    ///
    /// # Parameters
    /// - `log_types`: Optional vector of log types to collect.
    /// - `nonce`: Optional nonce for challenge-response.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// - The Ascend NPU API configuration is not initialized.
    /// - Failed to collect the specified log types.
    ///
    /// # Returns
    /// A `Result` containing the HashMap of log types and their corresponding AscendNpuEvidence or a `PluginError`.
    fn collect_log(&self, log_types: Option<Vec<String>>, nonce: &mut Vec<u8>) -> Result<HashMap<String, AscendNpuEvidence>, PluginError> {
        let mut mask: u64 = EVIDENCE_TYPE_EVIDENCE_TYPE_QUOTE;
        match log_types {
            Some(types) => {
                for item in &types {
                    match item.as_str() {
                        ASCEND_NPU_IMA_LOG => {
                            if self.config.enable_runtime_measurements {
                                mask += EVIDENCE_TYPE_EVIDENCE_TYPE_RUNTIME_MEASUREMENTS;
                            } else {
                                return Err(PluginError::InputError(
                                    "Runtime measurements is disabled".to_string(),
                                ));
                            }
                        }
                        ASCEND_NPU_TCG_EVENT_LOG => {
                            if self.config.enable_boot_measurements {
                                mask += EVIDENCE_TYPE_EVIDENCE_TYPE_BOOT_MEASUREMENTS;
                            } else {
                                return Err(PluginError::InputError(
                                    "Boot measurements is disabled".to_string(),
                                ));
                            }
                        }
                        other => {
                            return Err(PluginError::InputError(format!("Invalid log type {}", other)));
                        }
                    }
                }
            }
            None => {
                if self.config.enable_boot_measurements {
                    mask += EVIDENCE_TYPE_EVIDENCE_TYPE_BOOT_MEASUREMENTS;
                }
                if self.config.enable_runtime_measurements {
                    mask += EVIDENCE_TYPE_EVIDENCE_TYPE_RUNTIME_MEASUREMENTS;
                }
            }
        };
        let map = self.get_evidence(mask, nonce)?;
        Ok(map)
    }

    /// # get_evidence
    ///
    /// Collects evidence from the Ascend NPU device for the specified log types.
    ///
    /// # Parameters
    /// - `mask`: Mask to specify the log types to collect.
    /// - `nonce`: Optional nonce for challenge-response.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// - The Ascend NPU API configuration is not initialized.
    /// - Failed to collect the specified log types.
    ///
    /// # Returns
    /// A `Result` containing the HashMap of log types and their corresponding AscendNpuEvidence or a `PluginError`.
    fn get_evidence(&self, mask: u64, nonce: &mut Vec<u8>) -> Result<HashMap<String, AscendNpuEvidence>, PluginError> {
        // api config
        let api = ASCEND_NPU_API_CONFIG.get().ok_or_else(|| PluginError::InternalError("Failed to get ASCEND_NPU_API_CONFIG".to_string()))?;
        // get card list param
        let mut card_num = 0;
        let mut card_list = [0; 64];
        // nonce
        let input_ptr = nonce.as_mut_ptr();
        let input_len = nonce.len() as u32;

        let mut map = HashMap::<String, AscendNpuEvidence>::new();
        let mut ret: std::os::raw::c_int;
        unsafe {
            ret = (api.dcmi_init)();
        }
        if ret != 0 {
            return Err(PluginError::InternalError(format!("Failed to init dcmi, ret: {}", ret)));
        }
        unsafe {
            ret = (api.dcmi_get_card_list)(&mut card_num, card_list.as_mut_ptr(), card_list.len() as ::std::os::raw::c_int);
        }
        if ret != 0 {
            return Err(PluginError::InternalError(format!("Failed to get card list, ret: {}", ret)));
        }
        let mut attest_ctx = AttestCtx {
            device_id: 0, // The current configuration has two candidate values, but the two candidate values are indistinguishable.
            card_id: 0,
            logic_device_id: 0,
        };
        for card_id in &card_list[0..card_num as usize] {
            attest_ctx.card_id = *card_id as u32;
            let mut evidence = vec![0u8; 1024 * 1024];
            let mut evidence_len = evidence.len() as u32;
            unsafe {
                ret = (api.dcmi_attest_get_evidence)(&mut attest_ctx as *mut AttestCtx, input_ptr, input_len, evidence.as_mut_ptr(), &mut evidence_len, mask);
            }
            if ret == OUT_OF_MEMORY {
                evidence.reserve((evidence_len) as usize);
                unsafe {
                    ret = (api.dcmi_attest_get_evidence)(&mut attest_ctx as *mut AttestCtx, input_ptr, input_len, evidence.as_mut_ptr(), &mut evidence_len, mask);
                }
            }
            if ret != 0 {
                return Err(PluginError::InternalError(format!("Failed to get ascend_npu evidence, ret: {}", ret)));
            }
            let mut ak_cert = vec![0u8; 1024 * 1024];
            let mut ak_cert_len = ak_cert.len() as u32;
            unsafe {
                ret = (api.dcmi_attest_get_akcert)(&mut attest_ctx as *mut AttestCtx, ak_cert.as_mut_ptr(), &mut ak_cert_len);
            }
            if ret != 0 {
                return Err(PluginError::InternalError(format!("Failed to get ascend_npu akcert, ret: {}", ret)));
            }
            ak_cert.truncate(ak_cert_len as usize);
            evidence.truncate(evidence_len as usize);
            let mut dcmi_evidence: DcmiEvidence = match serde_cbor::from_slice(&evidence) {
                Ok(ev) => ev,
                Err(e) => return Err(PluginError::InternalError(format!("Failed to deserialize dcmi evidence: {}", e))),
            };
            dcmi_evidence.ak_cert = ak_cert;
            let evidence_msg = AscendNpuEvidence::try_from(dcmi_evidence)?;
            map.insert(card_id.to_string(), evidence_msg);
        }
        Ok(map)
    }
}

#[no_mangle]
pub fn create_plugin(host_functions: &AgentHostFunctions, plugin_type: &str) -> Result<Box<dyn AgentPlugin>, Box<dyn Error>> {
    let query_configuration = host_functions.query_configuration;
    AscendNpuPlugin::new(String::from(plugin_type), query_configuration)
        .map(|plugin| Box::new(plugin) as Box<dyn AgentPlugin>)
        .map_err(|e| Box::new(e) as Box<dyn Error>)
}